package main

import (
	"io"
	"log"
	"net"
	"net/http"

	"github.com/davecgh/go-spew/spew"
	"github.com/gorilla/websocket"
)

// 升级器配置
var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

func main() {
	http.HandleFunc("/", handleWebSocket)

	log.Println("WebSocket 服务启动在 :20001")
	log.Fatal(http.ListenAndServe("127.0.0.1:20001", nil))
}

func handleWebSocket(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("升级到 WebSocket 失败:", err)
		return
	}
	defer conn.Close()

	tcpConn, err := net.Dial("tcp", "127.0.0.1:20000")
	if err != nil {
		log.Println("连接到 TCP 服务失败:", err)
		return
	}
	defer tcpConn.Close()

	go func() {
		for {
			messageType, message, err := conn.ReadMessage()
			spew.Dump(message, messageType, err)
			if err != nil {
				log.Println("读取 WebSocket 消息失败:", err)
				return
			}
			if _, err := tcpConn.Write(message); err != nil {
				log.Println("写入 TCP 连接失败:", err)
				return
			}
			if messageType == websocket.CloseMessage {
				return
			}
		}
	}()

	buffer := make([]byte, 1024)
	for {
		n, err := tcpConn.Read(buffer)
		if n > 0 { // 确保有数据才写回 WebSocket
			if writeErr := conn.WriteMessage(websocket.TextMessage, buffer[:n]); writeErr != nil && writeErr != io.EOF {
				log.Println("写入 WebSocket 连接失败:", writeErr)
				return
			}
		}

		if n == 0 || (err != nil && err != io.EOF) { // 检查 EOF 和其他错误情况
			log.Println("读取 TCP 连接失败或连接关闭:", err)
			return
		}
	}
}
